#!/usr/bin/env python

# Entension of tutorial bot that randomly explores new territory
#
# Author: Bennett Wineholt
#
# Collaborative Diffusion chosen from reading Benjamin Meyer
#     with modifications envisioned by Gautam Kamath

import pickle
import Queue
import sys

from copy import deepcopy
from random import shuffle
from scipy import *
from scipy import ndimage

from ants import *


FRONTIER_LIMIT = 10
FRONTIER_SPARSIFY = 5
SHUFFLE_FRONTIER = False
VISION_TURNOVER = 100

TRANSITION = 0.125
REMAIN = 0.5

INFLUENCE_SPREAD = 10
FOOD_INFLUENCE_SPREAD = INFLUENCE_SPREAD
FRONTIER_INFLUENCE_SPREAD = 30
MY_ANTHILL_INFLUENCE_SPREAD = 4

FOOD_INFLUENCE = (500/TRANSITION)**FOOD_INFLUENCE_SPREAD
FRONTIER_INFLUENCE = (10/TRANSITION)**FRONTIER_INFLUENCE_SPREAD
#FRIENDLY_INFLUENCE = 0#(-1/TRANSITION)**INFLUENCE_SPREAD
MY_ANTHILL_INFLUENCE = 0#(-1/TRANSITION)**MY_ANTHILL_INFLUENCE_SPREAD
#ENEMY_INFLUENCE = 0#(-1/TRANSITION)**INFLUENCE_SPREAD
#ENEMY_ANTHILL_INFLUENCE = (1000/TRANSITION)**INFLUENCE_SPREAD
NEGLIGIBLE_INFLUENCE = 0.0001


class MyBot:
    def __init__(self):
        pass
 
    def do_setup(self, ants):
        #ENEMY_ANTHILL_SPREAD = max(ants.rows, ants.cols)
        #FOOD_SPREAD = sqrt(int(ants.viewradius2))
        #FRONTIER_SPREAD = ants.viewradius2
        #ENEMY_SPREAD = int(sqrt(ants.attackradius2))+1
        # pixels never seen
        self.unseen = []
        for row in range(ants.rows):
            for col in range(ants.cols):
                self.unseen.append((row, col))
        self.seen = Queue.Queue()
        # initialize one ant step filter
        self.one_step = array([[0,TRANSITION,0],
                               [TRANSITION,REMAIN,TRANSITION],
                               [0,TRANSITION,0]])
        # land pixels adjacent to pixels never seen
        self.frontier = None
        self.hills = [] 
        # total influence score of each map pixel
        self.influence_map = zeros((ants.rows, ants.cols))
        # steadily reseed influence map for convexity
        self.influence_seed = zeros((ants.rows, ants.cols))
        # influence inflation correction for water causes nodes to transmit
        #     (give) the proper value to nearby neighbors, then a
        #     self-influence correction is applied
        self.inflation = zeros((ants.rows, ants.cols))
    
    def do_turn(self, ants):

        # Begin helper method declarations

        def do_move_direction(loc, direction):
            new_loc = ants.destination(loc, direction)
            if (ants.unoccupied(new_loc) and new_loc not in orders):
                ants.issue_order((loc, direction))
                orders[new_loc] = loc
                targets[new_loc] = loc
                return True
            else:
                return False

        def do_move_location(loc, dest):
            directions = ants.direction(loc, dest)
            for direction in directions:
                if do_move_direction(loc, direction):
                    targets[dest] = loc
                    return True
            return False

        # Returns complement of known water on map
        def find_not_water():
            # matrix size of map marking pixels that are not water
            return ((array(ants.map) - WATER) != 0).astype(int)

        # Find land pixels adjacent to pixels never seen
        #
        # Returns tuple of
        #  - list of tuples   [(row, col), ... ]
        #  - frontier map - binary matrix of frontier pixels
        def find_frontier():
            M = array(ants.map)
            # matrix size of map marking land
            MLF = ((M - LAND*ones(M.shape)) == 0).astype(int)
            # matrix size of map marking unseen
            MUF = ((M - UNSEEN*ones(M.shape)) == 0).astype(int)
            # matrix size of map with unseen blurred one pixel
            F = array([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
            MUFF = ndimage.correlate(MUF, F)

            # intersect land and blurred unseen territory
            frontier_map = MLF & MUFF
            (rows, cols) = nonzero(frontier_map)

            res = [(rows[i], cols[i]) for i in range(0, len(rows))]
            return (res, frontier_map)

        # calculate inflation for influence to correct for loss by water sink
        #
        # Requires up to date self.not_water, isolated land pixels cause
        # warning to be generated, but will not corrupt influence map
        #
        # Returns tuple
        # - influence inflation correction for each pixel node
        # - number of water neighbors for each pixel node
        def calculate_inflation():
            assert (sum(self.not_water==1)+sum(self.not_water==0) ==
                    ants.rows*ants.cols)
            lost = (ndimage.correlate((1 - self.not_water).astype(float64),
                        array([[0,1,0],[1,0,1],[0,1,0]]), mode='wrap') *
                            self.not_water)
            assert (sum(lost==4) == 0)
            return 4.0 / ((4-lost).astype(float64))

        # propagates influence along land by one step
        #
        # inflation corrects for the loss of value given to neighbors since
        # some neighbors are water, 0 sink
        #
        # similarly, to utilize desirable properties of the normalized kernel
        # we re-adjust for the inflated self-influence and maintain balance
        #
        # Requires all scipy array
        #
        # Returns updated map
        def propagate_influence(imap):
            return ((ndimage.correlate(imap, self.one_step)) *
                    self.not_water)
            #return ((ndimage.correlate((imap*self.inflation) +
            #         self.influence_seed, self.one_step) -
            #         REMAIN * (self.inflation-1)) *
            #        self.not_water)

        def process_influence(points, influence, spread):
            imap = zeros((ants.rows, ants.cols))
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = influence
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return imap

        def process_influence_plus_euclid2(points, influence, spread, origin):
            imap = zeros((ants.rows, ants.cols))
            (o_r, o_c) = origin
            for (pt_r, pt_c) in points:
                imap[pt_r][pt_c] = (influence +
                                    ((o_r-pt_r)**2 + (o_c-pt_c)**2))
            for i in range(0, spread):
                imap = propagate_influence(imap)
            return imap

        # update the influence map for a turn with values of each location
        # based on
        # - food
        # - frontier (first FRONTIER_LIMIT pixels listed)
	# ##- friendly ants
        #        
        #
        def process_influence_map():
            imap = zeros((ants.rows, ants.cols))
            imap = imap + process_influence(ants.food(),
                                            FOOD_INFLUENCE,
                                            FOOD_INFLUENCE_SPREAD)
            imap = imap + process_influence_plus_euclid2(
                              self.frontier[:FRONTIER_LIMIT],
                              FRONTIER_INFLUENCE,
                              FRONTIER_INFLUENCE_SPREAD,
                              ants.my_hills()[0])
            imap = imap + process_influence(ants.my_hills(),
                                            MY_ANTHILL_INFLUENCE,
                                            MY_ANTHILL_INFLUENCE_SPREAD)
            #imap = imap + process_influence(ants.enemy_hills,
            #                                ENEMY_ANTHILL_INFLUENCE,
            #                                ENEMY_ANTHILL_INFLUENCE_SPREAD)       
            return imap

        # End helper method declarations


        # orders to return for each ant
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        orders = {}
        # destination targets to prevent conflicting pursuit of food
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        targets = {}

        # update visible map locations
        #if mod(ants.turns, VISION_RESET) == 0:
        #    self.unseen = []
        #    for row in range(ants.rows):
        #        for col in range(ants.cols):
        #            self.unseen.append((row, col))
        for i in range(0, min(VISION_TURNOVER, self.seen.qsize())):
            self.unseen.append(self.seen.get())
        for loc in self.unseen[:]:
            if ants.visible(loc):
                self.seen.put(loc)
                self.unseen.remove(loc)

        # update obstacle map, aka what's not water
        self.not_water = find_not_water()

        # get new frontier
        # frontier is land adjacent to unseen pixels
        (self.frontier, self.frontier_map) = find_frontier()
        #sparse_factor = max(int(len(self.frontier)/30), 1)
        self.frontier = self.frontier[::FRONTIER_SPARSIFY]
        if SHUFFLE_FRONTIER:
            shuffle(self.frontier)

        # calculate influence inflation correction
        self.inflation = calculate_inflation()

        # influence map accumulates value of map locations
        self.influence_map = process_influence_map()

        # prevent stepping on own hill
        for hill_loc in ants.my_hills():
            orders[hill_loc] = None

        # find close food
        #ant_dist = []
        #food_list = ants.food()
        #shuffle(food_list)
        #for food_loc in food_list:
        #    for ant_loc in ants.my_ants():
        #        dist = ants.distance(ant_loc, food_loc)
        #        ant_dist.append((dist, ant_loc, food_loc))
        #ant_dist.sort()
        #for dist, ant_loc, food_loc in ant_dist:
        #    if food_loc not in targets and ant_loc not in targets.values():
        #        do_move_location(ant_loc, food_loc)

        # find food and explore based on collaborative diffusion
        for (ant_r, ant_c) in ants.my_ants():
            pos_dest = []
            for direction in list(AIM.keys()):
                (new_r, new_c) = ants.destination((ant_r, ant_c), direction)
                pos_dest.append((self.influence_map[new_r][new_c],
                                 (new_r, new_c), direction))
            #pos_dest.append((self.influence_map[ant_r][ant_c], (ant_r, ant_c), ''))
            list.sort(pos_dest, reverse=True)
            #if abs(pos_dest[0][0]) < NEGLIGIBLE_INFLUENCE:
            #    break
            for (score, dest, direction) in pos_dest:
                if (dest not in targets and 
                        (ant_r, ant_c) not in targets.values() and
                        ants.unoccupied(dest)):
                    #if direction == '':
                    #    break
                    if do_move_direction((ant_r, ant_c), direction):
                        break

        # attack hills
        for hill_loc, hill_owner in ants.enemy_hills():
            if hill_loc not in self.hills:
                self.hills.append(hill_loc)        
        ant_dist = []
        for hill_loc in self.hills:
            for ant_loc in ants.my_ants():
                if ant_loc not in orders.values():
                    dist = ants.distance(ant_loc, hill_loc)
                    ant_dist.append((dist, ant_loc, hill_loc))
        ant_dist.sort()
        for dist, ant_loc, hill_loc in ant_dist:
            do_move_location(ant_loc, hill_loc)

        '''
        # explore frontier
        ant_dist = []
        for loc in self.frontier:
            for ant_loc in ants.my_ants():
                if ant_loc not in orders.values():
                    dist = ants.distance(ant_loc, loc)
                    ant_dist.append((dist, ant_loc, loc))
        ant_dist.sort()
        for dist, ant_loc, loc in ant_dist:
            do_move_location(ant_loc, loc) 
        '''

        
        # random move
        # has same effect as exploring unseen without worrying about pathing
        #
        # TODO(bmw): is 2 u-sec enough to terminate?
        for ant_loc in ants.my_ants():
            if ant_loc not in orders.values(): 
                if (ants.time_remaining() < 2):
                    break
                directions = list(AIM.keys())
                shuffle(directions)
                for direction in directions:
                    dest = ants.destination(ant_loc, direction)
                    if (not dest in targets and
                            ants.unoccupied(dest)):
                        if do_move_location(ant_loc, dest):
                            break
        

        '''
        # explore unseen areas
        for ant_loc in ants.my_ants():
            if ant_loc not in orders.values():
                unseen_dist = []
                for unseen_loc in self.unseen:
                    dist = ants.distance(ant_loc, unseen_loc)
                    unseen_dist.append((dist, unseen_loc))
                unseen_dist.sort()
                for dist, unseen_loc in unseen_dist:
                    if do_move_location(ant_loc, unseen_loc):
                        break
        '''

        # unblock own hill
        for hill_loc in ants.my_hills():
            if hill_loc in ants.my_ants() and hill_loc not in orders.values():
                for direction in ('s','e','w','n'):
                    if do_move_direction(hill_loc, direction):
                        break

        
if __name__ == '__main__':
    # psyco will speed up python a little, but is not needed
    try:
        import psyco
        psyco.full()
    except ImportError:
        pass
    
    try:
        # if run is passed a class with a do_turn method, it will do the work
        # this is not needed, in which case you will need to write your own
        # parsing function and your own game state class
        Ants.run(MyBot())
    except KeyboardInterrupt:
        print('ctrl-c, leaving ...')
